# coding: UTF-8
import argparse
import os
import time
from importlib import import_module

import numpy as np
import torch
from torch.utils.data import DataLoader

from dataset.Meta_dataset import MetaDataset
from loss.FocalLoss import FocalLoss
from train_eval.train_eval import train, test
from utils.utils import get_time_dif

if __name__ == '__main__':
    dataset = 'data/' + 'total7'
    print('Training on dataset:', dataset)
    model_name = 'ERNIE'
    x = import_module('models.' + model_name)
    config1 = (x.Config(dataset))
    model = x.Model(config1).to(config1.device)
    print('finish 1')
    y = import_module('MAML.' + 'meta')
    config2 = (y.Config())
    print('finish 2')
    maml = y.Meta(config2, model).to(config1.device)
    task_num_total = 100
    metadataset = MetaDataset(config1, config2, task_num_total, 'train')
    metadataset_test2 = MetaDataset(config1, config2, 5, 'test2')
    metadataset_test3 = MetaDataset(config1, config2, 5, 'test3')
    for i in range(task_num_total // config2.task_per_batch):
        a, b, c, d = metadataset.getbatch(i)
        accs = maml(a, b, c, d)
    accs_all_test = []
    for i in range(5):
        a, b, c, d = metadataset_test2.gettask(i)
        accs = maml.finetunning(a, b, c, d)
        accs_all_test.append(accs)
        print('Test acc:', accs)
    accs_all_test = []
    for i in range(5):
        a, b, c, d = metadataset_test3.gettask(i)
        accs = maml.finetunning(a, b, c, d)
        accs_all_test.append(accs)
        print('Test acc:', accs)

    